from ultralytics import YOLO
import cv2
import numpy as np
from shapely.geometry import Point, Polygon



# ============ 消防通道占用模型 ============
# 消防通道占用模型权重（根据 config.yaml）
weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\服务器上的权重\fire_lane_blockage.pt"
# 测试视频路径（换成您实际的消防通道视频）
video_path = r"E:\zhihao\2025文档\25年10月\Captures\消防通道\消防通道_1.mp4"


# ============ 获取视频信息 ============
cap = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
cap.release()


# ============ 运行消防通道占用检测 ============
model = YOLO(weight_path)

model.predict(
    video_path,
    vid_stride=1,                    # 每帧都检测
    imgsz=(height, width),           # 使用原始分辨率
    classes=None,                    # 检测所有类别（人、车、物体等）
    save=True,                       # 保存结果视频
    show=True,                       # 实时显示
    show_labels=True,                # 显示标签
    show_conf=True,                  # 显示置信度
    conf=0.3,                        # 置信度阈值 0.3
    iou=0.5                         # IOU 阈值
)

print("检测完成！")